import yaml
import threading
import time
import json
import os
import socket
import select
from collections import defaultdict
from threading import Lock
from datetime import datetime, timedelta
import sys

class PythonRedis:
    def __init__(self, config_path='redis_config.yaml'):
        # 加载配置
        self.load_config(config_path)
        
        # 初始化数据库存储
        self.databases = [{} for _ in range(16)]  # 16个数据库
        self.expiry_times = [{} for _ in range(16)]
        self.type_store = [{} for _ in range(16)]
        self.data_locks = [Lock() for _ in range(16)]
        
        # 当前选中的数据库
        self.current_db = 0
        
        # 启动清理过期键的后台任务
        self.start_expiry_checker()
        
        # 启动数据持久化任务
        self.start_persistence_task()
        
        # 启动网络服务
        self.start_server()

    def load_config(self, config_path):
        """加载配置文件"""
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
            
    def set(self, key, value, ex=None):
        """设置键值对，支持过期时间"""
        with self.data_locks[self.current_db]:
            self.databases[self.current_db][key] = value
            if ex is not None:
                self.expiry_times[self.current_db][key] = datetime.now() + timedelta(seconds=ex)
            # 保存到磁盘
            self.save_to_disk()
            return True

    def get(self, key):
        """获取值"""
        with self.data_locks[self.current_db]:
            if key in self.databases[self.current_db]:
                if self._is_expired(key):
                    self._delete_key(key)
                    return None
                return self.databases[self.current_db][key]
            return None

    def delete(self, key):
        """删除键"""
        with self.data_locks[self.current_db]:
            if key in self.databases[self.current_db]:
                self._delete_key(key)
                return True
            return False

    def _delete_key(self, key):
        """内部删除键的方法"""
        del self.databases[self.current_db][key]
        if key in self.expiry_times[self.current_db]:
            del self.expiry_times[self.current_db][key]

    def _is_expired(self, key):
        """检查键是否过期"""
        if key in self.expiry_times[self.current_db]:
            return datetime.now() > self.expiry_times[self.current_db][key]
        return False

    # Hash 操作
    def hset(self, name, key, value):
        """哈希表设置"""
        with self.data_locks[self.current_db]:
            if name not in self.databases[self.current_db]:
                self.databases[self.current_db][name] = {}
                self.type_store[self.current_db][name] = 'hash'
            self.databases[self.current_db][name][key] = value
            return True

    def hget(self, name, key):
        """哈希表获取"""
        with self.data_locks[self.current_db]:
            if name in self.databases[self.current_db] and isinstance(self.databases[self.current_db][name], dict):
                return self.databases[self.current_db][name].get(key)
            return None

    # List 操作
    def lpush(self, name, *values):
        """列表左推入"""
        with self.data_locks[self.current_db]:
            if name not in self.databases[self.current_db]:
                self.databases[self.current_db][name] = []
                self.type_store[self.current_db][name] = 'list'
            self.databases[self.current_db][name] = list(values) + self.databases[self.current_db][name]
            return len(self.databases[self.current_db][name])

    def rpop(self, name):
        """列表右弹出"""
        with self.data_locks[self.current_db]:
            if name in self.databases[self.current_db] and isinstance(self.databases[self.current_db][name], list):
                if len(self.databases[self.current_db][name]) > 0:
                    return self.databases[self.current_db][name].pop()
            return None

    # Set 操作
    def sadd(self, name, *values):
        """集合添加"""
        with self.data_locks[self.current_db]:
            if name not in self.databases[self.current_db]:
                self.databases[self.current_db][name] = set()
                self.type_store[self.current_db][name] = 'set'
            self.databases[self.current_db][name].update(values)
            return len(values)

    def smembers(self, name):
        """获取集合所有成员"""
        with self.data_locks[self.current_db]:
            if name in self.databases[self.current_db] and isinstance(self.databases[self.current_db][name], set):
                return list(self.databases[self.current_db][name])
            return []

    def start_expiry_checker(self):
        """启动过期键检查器"""
        def check_expired_keys():
            while True:
                with self.data_locks[self.current_db]:
                    now = datetime.now()
                    expired_keys = [
                        key for key, expiry in self.expiry_times[self.current_db].items()
                        if now > expiry
                    ]
                    for key in expired_keys:
                        self._delete_key(key)
                time.sleep(1)

        thread = threading.Thread(target=check_expired_keys, daemon=True)
        thread.start()

    def save_to_disk(self):
        """保存数据到磁盘"""
        data_to_save = {
            'databases': self.databases,
            'expiry_times': [
                {k: v.isoformat() for k, v in db.items()}
                for db in self.expiry_times
            ],
            'type_store': self.type_store
        }
        
        data_dir = self.config['storage']['data_dir']
        os.makedirs(data_dir, exist_ok=True)
        
        with open(os.path.join(data_dir, 'dump.json'), 'w') as f:
            json.dump(data_to_save, f)

    def load_from_disk(self):
        """从磁盘加载数据"""
        data_dir = self.config['storage']['data_dir']
        dump_file = os.path.join(data_dir, 'dump.json')
        
        if os.path.exists(dump_file):
            with open(dump_file, 'r') as f:
                data = json.load(f)
                
                # 处理旧格式的数据文件
                if 'data' in data:
                    # 旧格式：只有一个数据库
                    self.databases[0] = data['data']
                    self.expiry_times[0] = {
                        k: datetime.fromisoformat(v)
                        for k, v in data['expiry_times'].items()
                    }
                    self.type_store[0] = data['type_store']
                else:
                    # 新格式：多数据库
                    self.databases = data['databases']
                    self.expiry_times = [
                        {k: datetime.fromisoformat(v) for k, v in db.items()}
                        for db in data['expiry_times']
                    ]
                    self.type_store = data['type_store']

    def start_persistence_task(self):
        """启动持久化任务"""
        def persist_data():
            while True:
                self.save_to_disk()
                time.sleep(self.config['storage']['snapshot_interval'])

        thread = threading.Thread(target=persist_data, daemon=True)
        thread.start()

    def start_server(self):
        """启动Redis服务器"""
        def handle_client(client_socket):
            authenticated = False
            try:
                while True:
                    # 读取客户端请求
                    data = client_socket.recv(1024)
                    if not data:
                        break
                    
                    # 解析Redis协议
                    command = self._parse_command(data)
                    if not command:
                        continue
                    
                    # 检查是否需要认证
                    if not authenticated and command[0].upper() != 'AUTH':
                        client_socket.send("-NOAUTH Authentication required.\r\n".encode())
                        continue
                    
                    # 执行命令
                    response = self._execute_command(command)
                    
                    # 如果是AUTH命令且成功，设置认证状态
                    if command[0].upper() == 'AUTH' and response == "+OK\r\n":
                        authenticated = True
                    
                    # 发送响应
                    client_socket.send(response.encode())
            except Exception as e:
                print(f"Error handling client: {e}")
            finally:
                client_socket.close()

        def server_loop():
            try:
                server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                
                # 从配置中获取主机和端口
                host = self.config['server']['host']
                port = self.config['server']['port']
                
                print(f"Attempting to start Redis server on {host}:{port}")
                sys.stdout.flush()
                
                try:
                    server.bind((host, port))
                    server.listen(5)
                    print(f"Redis server successfully started on {host}:{port}")
                    sys.stdout.flush()
                except OSError as e:
                    print(f"Failed to start server: {e}")
                    if e.errno == 10048:  # 端口被占用
                        print("Port is already in use. Please check if another Redis server is running.")
                    elif e.errno == 10049:  # 无法绑定地址
                        print("Cannot bind to address. Please check your network configuration.")
                    raise
                
                while True:
                    try:
                        client_socket, addr = server.accept()
                        print(f"New connection from {addr}")
                        thread = threading.Thread(target=handle_client, args=(client_socket,))
                        thread.daemon = True
                        thread.start()
                    except Exception as e:
                        print(f"Error accepting connection: {e}")
                        continue
                        
            except Exception as e:
                print(f"Server error: {e}")
                raise

        thread = threading.Thread(target=server_loop)
        thread.daemon = True
        thread.start()

    def _parse_command(self, data):
        """解析Redis协议命令"""
        try:
            parts = data.decode().split('\r\n')
            if not parts or parts[0][0] != '*':
                return None
            
            command_count = int(parts[0][1:])
            command = []
            
            for i in range(1, command_count * 2 + 1, 2):
                if parts[i][0] == '$':
                    command.append(parts[i+1])
            
            return command
        except:
            return None

    def _execute_command(self, command):
        """执行Redis命令并返回响应"""
        if not command:
            return "-ERR invalid command\r\n"
        
        cmd = command[0].upper()
        try:
            if cmd == 'AUTH':
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'auth' command\r\n"
                if command[1] == self.config['server']['password']:
                    return "+OK\r\n"
                return "-ERR invalid password\r\n"
            elif cmd == 'SELECT':
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'select' command\r\n"
                try:
                    db_index = int(command[1])
                    if 0 <= db_index < 16:
                        self.current_db = db_index
                        return "+OK\r\n"
                    return "-ERR DB index is out of range\r\n"
                except ValueError:
                    return "-ERR invalid DB index\r\n"
            elif cmd == 'SET':
                if len(command) < 3:
                    return "-ERR wrong number of arguments for 'set' command\r\n"
                with self.data_locks[self.current_db]:
                    self.databases[self.current_db][command[1]] = command[2]
                    self.save_to_disk()
                return "+OK\r\n"
            elif cmd == 'GET':
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'get' command\r\n"
                with self.data_locks[self.current_db]:
                    value = self.databases[self.current_db].get(command[1])
                    if value is None:
                        return "$-1\r\n"
                    return f"${len(str(value))}\r\n{value}\r\n"
            elif cmd == 'PING':
                return "+PONG\r\n"
            elif cmd == 'SUBSCRIBE':
                # 订阅命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'subscribe' command\r\n"
                channel = command[1]
                return f"*3\r\n$9\r\nsubscribe\r\n${len(channel)}\r\n{channel}\r\n:1\r\n"
            elif cmd == 'PSUBSCRIBE':
                # 模式订阅命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'psubscribe' command\r\n"
                pattern = command[1]
                return f"*3\r\n$10\r\npsubscribe\r\n${len(pattern)}\r\n{pattern}\r\n:1\r\n"
            elif cmd == 'UNSUBSCRIBE':
                # 取消订阅命令
                return "*3\r\n$11\r\nunsubscribe\r\n$1\r\n1\r\n:0\r\n"
            elif cmd == 'PUNSUBSCRIBE':
                # 取消模式订阅命令
                return "*3\r\n$12\r\npunsubscribe\r\n$1\r\n1\r\n:0\r\n"
            elif cmd == 'PUBLISH':
                # 发布命令
                if len(command) < 3:
                    return "-ERR wrong number of arguments for 'publish' command\r\n"
                return ":1\r\n"  # 返回订阅者数量
            elif cmd == 'CONFIG':
                # CONFIG命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'config' command\r\n"
                if command[1].upper() == 'GET':
                    if len(command) < 3:
                        return "-ERR wrong number of arguments for 'config get' command\r\n"
                    param = command[2]
                    if param == "timeout":
                        return "*2\r\n$7\r\ntimeout\r\n$1\r\n0\r\n"
                    elif param == "notify-keyspace-events":
                        return "*2\r\n$19\r\nnotify-keyspace-events\r\n$1\r\nA\r\n"
                    return "*0\r\n"
                return "+OK\r\n"
            elif cmd == 'CLIENT':
                # CLIENT命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'client' command\r\n"
                subcmd = command[1].upper()
                if subcmd == 'SETNAME':
                    return "+OK\r\n"
                elif subcmd == 'GETNAME':
                    return "$-1\r\n"
                elif subcmd == 'LIST':
                    return "*0\r\n"
                return "+OK\r\n"
            elif cmd == 'INFO':
                # INFO命令
                return "$0\r\n\r\n"
            elif cmd == 'QUIT':
                # QUIT命令
                return "+OK\r\n"
            elif cmd == 'TIME':
                # TIME命令
                now = datetime.now()
                return f"*2\r\n${len(str(int(now.timestamp())))}\r\n{int(now.timestamp())}\r\n${len(str(now.microsecond))}\r\n{now.microsecond}\r\n"
            elif cmd == 'ECHO':
                # ECHO命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'echo' command\r\n"
                return f"${len(command[1])}\r\n{command[1]}\r\n"
            elif cmd == 'EXISTS':
                # EXISTS命令
                if len(command) < 2:
                    return "-ERR wrong number of arguments for 'exists' command\r\n"
                with self.data_locks[self.current_db]:
                    exists = command[1] in self.databases[self.current_db]
                    return f":{1 if exists else 0}\r\n"
            else:
                print(f"Unknown command: {cmd}")
                return "-ERR unknown command\r\n"
        except Exception as e:
            print(f"Command execution error: {e}")
            return f"-ERR {str(e)}\r\n"

if __name__ == "__main__":
    try:
        # 创建Redis实例
        redis = PythonRedis('redis_config.yaml')
        print("Redis server initialized successfully")
        
        # 保持程序运行
        while True:
            time.sleep(1)
    except Exception as e:
        print(f"Failed to start Redis server: {e}")
        sys.exit(1)
